# coding: utf-8
"""
使用说明:
1、excel文件格式(参见"7分支.xls")
(1) 必须有一个"分支"子表
(2) "分支"表中总共4列数据, 前3列拓扑关系, 第4列q表示已知风量(未知的不用填)
(3) 通风网络必须是无源汇网络, 必须确保每个节点都有入边和出边, 这样才能保证每个节点满足流量守恒
    注: 通过手动增加虚拟大气分支的方法
2、已知风量的个数: n-m+1
   注: 这里的n表示总分支个数, 包括虚拟大气分支
"""
from jl import (
    read_xls_file as _read_xls_file,
    build_network as _build_network,
)
from functools import cmp_to_key
from igraph import Graph, OUT, IN
import numpy as np

# 读取xls文件, 构造有向图对象
def build_graph(xls_file):
    json_from_client = _read_xls_file(xls_file)
    # print("读取xls文件", xls_file, "生成json:", json_from_client)
    dg = Graph(directed=True)
    _build_network(json_from_client, dg)
    return dg


def custom_sort(v1, v2):
    id1, id2 = int(v1["id"]), int(v2["id"])
    if id1 > id2:
        return 1
    elif id1 < id2:
        return -1
    else:
        return 0


def build_equation(dg, edge_cols, in_edges, out_edges):
    a = [0] * len(edge_cols)
    # 查找节点的入边(+1)
    for i in in_edges:
        j = edge_cols[i]
        a[j] = 1
    # 查找节点的出边(-1)
    for i in out_edges:
        j = edge_cols[i]
        a[j] = -1
    return a


# 构造邻接矩阵
def build_adjacent_matrix(dg, edge_cols, nodes):
    A = []
    # 遍历每个节点构造邻接矩阵
    for u in nodes:
        # 忽略源点和汇点
        if u.degree(mode=OUT) == 0 or u.degree(mode=IN) == 0:
            continue
        _a = build_equation(
            dg, edge_cols, dg.incident(u, mode=IN), dg.incident(u, mode=OUT),
        )
        A.append(_a)

    # # 源点和汇点再列一个方程
    # # 查找汇点
    # t_nodes = [u for u in dg.vs.select(_outdegree=0)]
    # # 查找源点
    # s_nodes = [u for u in dg.vs.select(_indegree=0)]
    # if len(s_nodes) > 0 and len(t_nodes) > 0:
    #     # 这里仅考虑1个源点和1个汇点
    #     s, t = s_nodes[0], t_nodes[0]
    #     _a = build_equation(
    #         dg, edge_cols, dg.incident(t, mode=IN), dg.incident(s, mode=OUT),
    #     )
    #     A.append(_a)
    # 任意去掉一行(这里为了方便，直接删除最后一行)
    A = np.delete(A, -1, axis=0)
    return A, edge_cols


# 构造线性方程组
def build_linear_equation(dg, A, edge_cols, known_edges):
    # 移动已知风量的到右边
    indexs = []
    _A = np.array(A)
    _b = np.zeros(_A.shape[0])

    # 记录每一列对应的分支id
    cols = [-1] * _A.shape[1]
    for i in edge_cols:
        j = edge_cols[i]
        cols[j] = i

    for i in known_edges:
        j = edge_cols[i]
        # 记录已知风量所在的列
        indexs.append(j)
        # 标记已知风量为-1
        cols[j] = -1 * i
        _b += -1 * dg.es[i]["q"] * _A[:, j]
    # 真实的b向量
    _b = _b.reshape(_b.shape[0], 1)
    # 删除邻接矩阵A中已知风量的列
    _A = np.delete(_A, indexs, axis=1)
    # print(_A)
    # print(_b)
    return _A, _b, [i for i in cols if i >= 0], [-1 * i for i in cols if i < 0]


# 求解矩阵
def solve_matrix(A, b):
    x = np.linalg.solve(A, b)
    # solve求得的x是一个二维数组, 维度(3,1) ==> x = [[2], [-1], [3]]
    # 需要变换为一维向量 => x = [2, -1, 3]
    x = x.reshape(x.shape[1], x.shape[0])[0]
    return x


def main():
    xls_file = "./7分支.xls"
    # 读取excel表格, 利用简化版的vno库构造有向图
    dg = build_graph(xls_file)
    # print(dg)
    # 将分支按
    edges = [e for e in dg.es]
    edge_cols = dict([(e.index, i) for i, e in enumerate(edges)])
    nodes = [u for u in dg.vs]
    # 确定已知的风量
    known_edges = [
        e.index for e in dg.es.select(lambda e: e["q"] != None and abs(e["q"]) > 0)
    ]
    # 构造邻接矩阵
    A, edge_cols = build_adjacent_matrix(dg, edge_cols, nodes)
    A, b, A_cols, b_cols = build_linear_equation(dg, A, edge_cols, known_edges)
    print("\n===================== 已知风量 =====================")
    print(
        "\t".join(["e%s = %.1f" % (str(dg.es[i]["id"]), dg.es[i]["q"]) for i in b_cols])
    )
    print("\n=============== 线性方程组: Ax = b =================")
    print("\t".join(["e%s" % (str(dg.es[i]["id"])) for i in A_cols]), "\t|\t  b")
    n, m = A.shape
    for i in range(n):
        for j in range(m):
            print(A[i, j], end="\t")
        print("|\t", b[i][0])

    # print(A_cols, b_cols)
    # 求解矩阵
    x = solve_matrix(A, b)
    # print("x=", x)
    # 合并所有的风量值(包括已知的风量和计算的风量)
    Q = [(A_cols[i], q) for i, q in enumerate(x)]
    Q.extend([(i, dg.es[i]["q"]) for i in b_cols])
    Q_dict = dict(Q)
    print("\n===================== 方程组解 =====================")
    n = 1
    for i in sorted(Q_dict.keys()):
        print("q%s = %.1f" % (str(dg.es[i]["id"]), Q_dict[i]), end="\t")
        if n % 3 == 0:
            print()
        n += 1
    # print(
    #     "\n".join(
    #         [
    #             "q%s = %.1f" % (str(dg.es[i]["id"]), Q_dict[i])
    #             for i in sorted(Q_dict.keys())
    #         ]
    #     )
    # )


if __name__ == "__main__":
    main()
